#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RRT双向模块

此模块提供双向RRT实现，它从起始位置和目标位置同时构建树，
以更高效地寻找路径。
"""

import numpy as np
import time

from .rrt_base import RRTBase
from .tree import Tree
from .node import Node


class RRTBidirectional(RRTBase):
    """
    双向RRT实现。
    
    此算法构建两棵树 - 一棵从起点，一棵从目标点，
    并尝试连接它们以比标准RRT更快地找到路径。
    """
    
    def __init__(self, collision_checker, config):
        """
        初始化双向RRT。
        
        参数:
            collision_checker: 碰撞检查器对象
            config (dict): 配置字典
        """
        super().__init__(collision_checker, config)
        
        # 双向特定参数
        self.bidir_config = config.get("bidirectional", {})
        self.connect_dist = self.bidir_config.get("connect_distance", 10.0)
        self.tree_swap_prob = self.bidir_config.get("tree_swap_probability", 0.5)
        
        # 第二棵树（目标树）
        self.goal_tree = None
        self.connection_node_start = None
        self.connection_node_goal = None
        
    def plan(self, start, goal, max_iterations=None, goal_tolerance=None, plotter=None):
        """
        使用双向RRT规划从起点到目标的路径。
        
        参数:
            start (tuple): 起始位置 (x, y)
            goal (tuple): 目标位置 (x, y)
            max_iterations (int, optional): 最大迭代次数
            goal_tolerance (float, optional): 目标达成容差
            plotter (Plotter, optional): 用于可视化的绘图器
            
        返回:
            tuple: (path, info) 其中path是点的列表，info是
                  包含额外信息的字典
        """
        # 如果提供，则覆盖参数
        if max_iterations is not None:
            self.max_iterations = max_iterations
        if goal_tolerance is not None:
            self.goal_tolerance = goal_tolerance
            
        # 初始化规划
        start_time = time.time()
        self._init_bidirectional_planning(start, goal)
        
        # 如果提供了绘图器，则初始化可视化
        if plotter:
            plotter.init_bidirectional_planning(self.tree, self.goal_tree, start, goal)
        
        # 主规划循环
        for i in range(self.max_iterations):
            # 更新迭代计数
            self.iterations = i + 1
            
            # 决定扩展哪棵树
            expand_start_tree = np.random.random() > self.tree_swap_prob
            
            if expand_start_tree:
                active_tree = self.tree
                target_tree = self.goal_tree
                target_point = goal
            else:
                active_tree = self.goal_tree
                target_tree = self.tree
                target_point = start
            
            # 采样随机点
            if np.random.random() < self.goal_bias:
                # 以goal_bias的概率采样目标点
                x_rand, y_rand = target_point
            else:
                # 以(1 - goal_bias)的概率采样随机点
                x_rand, y_rand = self._sample_free()
            
            # 尝试扩展活动树
            new_node = self._extend(active_tree, (x_rand, y_rand))
            
            # 如果扩展成功，尝试连接到另一棵树
            if new_node:
                # 尝试连接两棵树
                if self._try_connect_trees(active_tree, target_tree, new_node):
                    # 连接成功，我们找到了一条路径
                    self.goal_reached = True
                    break
            
            # 如果提供了绘图器，则更新可视化
            if plotter and i % 10 == 0:  # 每10次迭代更新一次
                plotter.update_bidirectional_planning(self.tree, self.goal_tree, 
                                                     self.connection_node_start, 
                                                     self.connection_node_goal)
        
        # 规划后提取路径
        self.planning_time = time.time() - start_time
        self.path = self._extract_bidirectional_path()
        
        # 如果提供了绘图器，则进行最终可视化更新
        if plotter:
            plotter.update_bidirectional_planning(self.tree, self.goal_tree, 
                                                 self.connection_node_start, 
                                                 self.connection_node_goal, self.path)
        
        # 返回路径和规划信息
        info = {
            "num_nodes_start_tree": self.tree.size,
            "num_nodes_goal_tree": self.goal_tree.size,
            "iterations": self.iterations,
            "planning_time": self.planning_time,
            "goal_reached": self.goal_reached,
            "start_tree_nodes": self.tree.nodes,
            "goal_tree_nodes": self.goal_tree.nodes,
            "connection_node_start": self.connection_node_start,
            "connection_node_goal": self.connection_node_goal
        }
        
        return self.path, info
    
    def _init_bidirectional_planning(self, start, goal):
        """
        初始化双向规划过程。
        
        参数:
            start (tuple): 起始位置 (x, y)
            goal (tuple): 目标位置 (x, y)
        """
        # 初始化起始树（与RRTBase中相同）
        self.start_node = Node(start[0], start[1])
        self.tree = Tree(self.start_node)
        
        # 初始化目标树
        self.goal_node = Node(goal[0], goal[1])
        self.goal_tree = Tree(self.goal_node)
        
        # 重置状态
        self.goal_reached = False
        self.iterations = 0
        self.path = None
        self.connection_node_start = None
        self.connection_node_goal = None
    
    def _extend(self, tree, point):
        """
        向给定点扩展树。
        
        参数:
            tree (Tree): 要扩展的树
            point (tuple): 要扩展的目标点
            
        返回:
            Node: 如果扩展成功则返回新节点，否则返回None
        """
        # 在树中找到最近的节点
        nearest_node = tree.get_nearest_node(point[0], point[1])
        
        # 朝随机点方向引导
        x_new, y_new = self._steer(nearest_node.position, point)
        
        # 检查到新点的路径是否无碰撞
        if not self._check_segment_collision(nearest_node.position, (x_new, y_new)):
            # 创建并将新节点添加到树中
            cost = nearest_node.cost + self._compute_distance(nearest_node.position, (x_new, y_new))
            new_node = Node(x_new, y_new, nearest_node, cost)
            tree.add_node(new_node)
            return new_node
        
        return None
    
    def _try_connect_trees(self, source_tree, target_tree, source_node):
        """
        尝试通过给定的源节点连接两棵树。
        
        参数:
            source_tree (Tree): 包含源节点的树
            target_tree (Tree): 要连接的目标树
            source_node (Node): 连接的源节点
            
        返回:
            bool: 如果树成功连接则为True，否则为False
        """
        # 在目标树中找到最近的节点
        target_node = target_tree.get_nearest_node(source_node.x, source_node.y)
        
        # 检查节点之间的距离
        distance = self._compute_distance(source_node.position, target_node.position)
        
        # 如果节点足够接近，尝试连接它们
        if distance <= self.connect_dist:
            # 检查节点之间是否有碰撞
            if not self._check_segment_collision(source_node.position, target_node.position):
                # 树已连接
                # 存储连接节点
                if source_tree == self.tree:
                    self.connection_node_start = source_node
                    self.connection_node_goal = target_node
                else:
                    self.connection_node_start = target_node
                    self.connection_node_goal = source_node
                
                return True
        
        return False
    
    def _extract_bidirectional_path(self):
        """
        双向规划后提取从起点到目标的路径。
        
        返回:
            list: 从起点到目标的(x, y)坐标列表
        """
        if not self.goal_reached:
            return []
            
        # 获取从起点到连接节点的路径
        path_from_start = self.tree.get_path_to_node(self.connection_node_start)
        
        # 获取从目标到连接节点的路径（反向）
        path_from_goal = self.goal_tree.get_path_to_node(self.connection_node_goal)
        path_from_goal.reverse()
        
        # 合并路径
        full_path = path_from_start + path_from_goal[1:]  # 跳过目标路径的第一个节点（重复）
        
        return full_path 